import os
import sys
import json
import jieba
import codecs
codecs.register_error("strict", codecs.replace_errors)


# import jieba


class MySentences(object):
    def __init__(self, dir_name):
        self.dir_name = dir_name

    def __iter__(self):
        for f_name in os.listdir(self.dir_name):
            for line in open(os.path.join(self.dir_name, f_name)):
                out = {}
                try:
                    out = json.loads(line)
                except Exception as e:
                    print(line)
                    print(str(e))
                yield out


if __name__ == '__main__':
    text_dir = '/data/data/category_lanmu/at'
    out_file1 = '/data/data/class1_title.tsv'
    out_file2 = '/data/data/class1_content.tsv'

    out_fd1 = open(out_file1, 'w+')
    out_fd2 = open(out_file2, 'w+')

    sentences = MySentences(text_dir)
    class1_dict = {}
    class2_dict = {}
    for s in sentences:
        class1 = s.get('article_type_1st', '')
        title = s.get('title', '').replace('\r\n', ' ').replace('\n', ' ')
        title = ' '.join(jieba.cut(title))
        content = s.get('content', '').replace('\r\n', ' ').replace('\n', ' ')
        content = ' '.join(jieba.cut(content))

        if class1:
            # out_fd1.write(class1 + '\t' + title + '\n')
            # out_fd2.write(class1 + '\t' + content + '\n')
            out_fd1.write(title + '__label__' + class1 + '\n')
            out_fd2.write(content + '__label__' + class1 + '\n')



